import copy
import arrow
import requests
import json

from pymongo import MongoClient

from utility.log import Logger
from conf.settings import CLOUDMONITOR_MONGODB, CLOUDMONITOR_DB, WASS_MONGODB, WASS_DB, SCANNER
from utility.enums import NotificationType

logger = Logger().get_logger()
replace_dict = {'hour': 1, 'minute': 0, 'second': 0, 'microsecond': 0}
arrow_now = arrow.now().replace(**replace_dict)
if arrow_now <= arrow.utcnow():
    arrow_now.shift(days=1)

db = MongoClient(CLOUDMONITOR_MONGODB)[CLOUDMONITOR_DB]
wass_db = MongoClient(WASS_MONGODB)[WASS_DB]


class Task(object):
    allow_keys = ('task_id', 'target_url', 'origin_ip', 'job_names', 'filter_events',
                  'task_config', 'job_config_map', 'flex_job_config_map', 'target_id', 'user_id', 'addition')

    url = SCANNER['base_url']

    def __init__(self, task_id):
        self.data = None
        self.task_id = task_id
        resp = requests.get(f'{self.url}/v1/tasks:get?task_id={task_id}')
        if resp.status_code != 200:
            logger.error(task_id)
            return
        origin_data = json.loads(resp.text)
        self.data = {k: origin_data.get(k) for k in self.allow_keys}

    def disable(self):
        if not self.data:
            return
        data = copy.deepcopy(self.data)
        data['task_config']['is_disabled'] = True

        for i in range(5):
            try:
                resp = requests.post(f"{self.url}/v1/tasks:modify", data=json.dumps(data))
                if resp.status_code != 200:
                    raise ValueError()
                break
            except:
                logger.error(f"{self.task_id} retry")


def get_scan_tasks() -> dict:
    query = {
        'task_config.is_periodic': True,
        'task_config.is_disabled': False,
        'addition.jobId': {'$exists': 1}
    }

    drop_key_tuple = ('_id', )
    job_scan_tasks_dict = {}
    for item in wass_db.scan_tasks.find(query):
        scan_task_list = job_scan_tasks_dict.setdefault(item['addition']['jobId'], [])
        for k in drop_key_tuple:
            item.pop(k)
        scan_task_list.append(item)

    return job_scan_tasks_dict


def get_disabled_scan_task_ids_from_job() -> list:
    disabled_task_type_dict = {}
    task_type_tuple = ('securityEvent', 'ssl')
    # asset_dict = {item['jobId']: item for item in db.assets.find()}
    for item in db.jobs.find():
        enable_monitor = item.get('enableMonitor', False)
        for task_type in task_type_tuple:
            if not (settings := item.get(f"{task_type}Settings")):
                continue
            if isinstance(settings, list):
                # TODO: http, ping
                continue
            else:
                if (task_id := settings.get('taskId')) and (not enable_monitor or not settings.get('enable', False)):
                    disabled_task_type_dict.setdefault(task_type, []).append(task_id)

    task_id_list = []
    for task_type, tid_list in disabled_task_type_dict.items():
        for item in wass_db.scan_tasks.find({'task_id': {'$in': tid_list}, 'task_config.is_disabled': False}):
            task_id_list.append(item.get('task_id'))

    return task_id_list


def get_disabled_task_id_list():
    task_id_list = []
    task_type_tuple = ('vul', 'ssl', 'securityEvent', 'content', 'http', 'ping')

    asset_dict = {item['jobId']: item for item in db.assets.find()}

    for item in db.jobs.find({'enableMonitor': True}):
        for task_type in task_type_tuple:
            if not (setting := item.get(f"{task_type}Settings")):
                continue
            if isinstance(setting, list):
                for st in setting:
                    if st.get('enable') and (_task_id := st.get('taskId')):
                        task_id_list.append(_task_id)
            else:
                if setting.get('enable') and (_task_id := setting.get('taskId')):
                    task_id_list.append(_task_id)

        asset = asset_dict.get(item['_id'])
        if not asset:
            continue
        setting = asset.get(f"taskSettings", {})
        if setting.get('enable') and (_task_id := setting.get('taskId')):
            task_id_list.append(_task_id)

    return task_id_list


def do_disable_scan_tasks_by_id(task_id_list):
    from queue import Queue
    from time import sleep
    from threading import Thread

    class T(Thread):
        def __init__(self, task_id, idx, q: Queue):
            self.task_id = task_id
            self.q = q
            self.idx = idx
            super(T, self).__init__()

        def run(self) -> None:
            try:
                Task(self.task_id).disable()
            finally:
                logger.info(self.idx)
                self.q.get()
                self.q.task_done()

    logger.info(f"Start {__file__}")

    q = Queue(maxsize=10)

    for index, item in enumerate(task_id_list):
        q.put(index)
        T(item, index, q).start()
        if index % 10 == 0:
            sleep(2)

    q.join()

    logger.info(f"End {__file__}")


def get_tasks_to_repeat(query):
    from utility.monitor import Monitor
    from queue import Queue
    from threading import Thread
    from copy import deepcopy

    class T(Thread):
        def __init__(self, t_id, idx, total, q: Queue):
            self.task_id_list = deepcopy(t_id)
            self.q = q
            self.idx = idx
            self.total = total
            super(T, self).__init__()

        def run(self) -> None:
            for x in range(3):
                try:
                    params = {
                        "notify_type_name": NotificationType.NT_ACTION_REPEAT.value,
                        "task": {
                            "addition": {
                                "task_id_list": self.task_id_list
                            }
                        }
                    }
                    resp = Monitor().post("/api/v2/notify", params)
                    if resp and resp.status_code == 200:
                        logger.info("Send to celery.")
                        logger.info(f"{self.idx}/{self.total}")
                        break
                finally:
                    self.q.get()
                    self.q.task_done()

    logger.info(f"Start {__file__}")

    task_id_list = []
    i = 0
    for task in db.tasks.find(query, {'_id': 1}):
        task_id_list.append(str(task.get('_id')))
        i += 1

    total = i

    logger.info(f"Find {total} items.")

    q = Queue(maxsize=100)

    cache = []
    for idx, _id in enumerate(task_id_list):
        cache.append(_id)
        if len(cache) >= 10:
            q.put(idx)
            T(cache, idx, total, q).start()
            cache = []

    if cache:
        q.put(total)
        T(cache, total, total, q).start()

    q.join()

    logger.info(f"End {__file__}")


def clear_tasks(query):
    """Remove invalid tasks from cloudmonitor"""

    logger.info(f"Start {__file__}")

    job_id_list = [str(item['_id']) for item in db.jobs.find({}, {'_id': 1})]

    remove_task_id_list = []
    for item in db.tasks.find(query, {'_id': 1, 'jobId': 1, 'target': 1, 'uid': 1}):
        if str(item['jobId']) not in job_id_list:
            remove_task_id_list.append(item['_id'])
            logger.info(f"Add task {{_id: {item['_id']}, target: {item['target']}, uid: {item['uid']}}} to remove list.")

    db.tasks.remove({'_id': {'$in': remove_task_id_list}})
    logger.info(f"Succeed in removing.")

    logger.info(f"End {__file__}")


def main():
    task_id_list = get_disabled_task_id_list()

    do_disable_scan_tasks_by_id(task_id_list)


if __name__ == '__main__':
    main()
